#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
系统根目录节点 - 获取系统所有根目录
"""

import os
import platform
import json
from typing import List, Dict, Any
import comfy.utils

class SystemRootsNode:
    """系统根目录节点 - 返回系统所有根目录"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "include_network": ("BOOLEAN", {"default": False}),
                "include_removable": ("BOOLEAN", {"default": True}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("root_paths", "root_info", "system_info")
    FUNCTION = "get_system_roots"
    CATEGORY = "AI训练/path_manager"

    def get_system_roots(self, include_network, include_removable):
        """获取系统所有根目录"""
        try:
            system = platform.system().lower()
            root_paths = []
            root_info = {}
            system_info = {
                "system": system,
                "platform": platform.platform(),
                "architecture": platform.architecture()[0]
            }
            
            if system == "windows":
                # Windows系统 - 获取所有盘符
                import string
                from ctypes import windll
                
                drives = []
                bitmask = windll.kernel32.GetLogicalDrives()
                
                for letter in string.ascii_uppercase:
                    if bitmask & 1:
                        drive = f"{letter}:\\"
                        drives.append(drive)
                    bitmask >>= 1
                
                # 获取每个驱动器的详细信息
                for drive in drives:
                    try:
                        drive_type = windll.kernel32.GetDriveTypeW(drive)
                        drive_info = {
                            "path": drive,
                            "type": self._get_drive_type_name(drive_type),
                            "drive_type_code": drive_type,
                            "exists": os.path.exists(drive),
                            "accessible": os.access(drive, os.R_OK)
                        }
                        
                        # 根据驱动器类型决定是否包含
                        should_include = True
                        if drive_type == 2:  # DRIVE_REMOVABLE
                            should_include = include_removable
                        elif drive_type == 4:  # DRIVE_REMOTE
                            should_include = include_network
                        
                        if should_include:
                            root_paths.append(drive)
                            root_info[drive] = drive_info
                            
                    except Exception as e:
                        print(f"获取驱动器 {drive} 信息失败: {e}")
            
            elif system == "linux":
                # 只显示 / 以及 / 下的所有一级文件夹
                root_paths = []
                root_info = {}
                try:
                    for entry in os.listdir("/"):
                        full_path = os.path.join("/", entry)
                        if os.path.isdir(full_path):
                            root_paths.append(full_path)
                            root_info[full_path] = {
                                "path": full_path,
                                "type": "dir",
                                "exists": True,
                                "accessible": os.access(full_path, os.R_OK)
                            }
                    # 把 / 也加进去
                    root_paths.insert(0, "/")
                    root_info["/"] = {
                        "path": "/",
                        "type": "root",
                        "exists": True,
                        "accessible": os.access("/", os.R_OK)
                    }
                except Exception as e:
                    print(f"读取 / 目录失败: {e}")
            
            elif system == "darwin":  # macOS
                # macOS系统 - 获取根目录和挂载点
                root_paths = ["/"]
                root_info["/"] = {
                    "path": "/",
                    "type": "root",
                    "exists": True,
                    "accessible": os.access("/", os.R_OK)
                }
                
                # 获取其他挂载点
                try:
                    import subprocess
                    result = subprocess.run(["mount"], capture_output=True, text=True)
                    if result.returncode == 0:
                        for line in result.stdout.split('\n'):
                            if line.strip():
                                parts = line.split()
                                if len(parts) >= 3:
                                    mount_point = parts[2]
                                    if mount_point not in root_paths and os.path.exists(mount_point):
                                        root_paths.append(mount_point)
                                        root_info[mount_point] = {
                                            "path": mount_point,
                                            "type": "mount",
                                            "exists": True,
                                            "accessible": os.access(mount_point, os.R_OK)
                                        }
                except Exception as e:
                    print(f"读取挂载点失败: {e}")
            
            else:
                # 其他系统 - 使用通用方法
                root_paths = ["/"]
                root_info["/"] = {
                    "path": "/",
                    "type": "root",
                    "exists": os.path.exists("/"),
                    "accessible": os.access("/", os.R_OK) if os.path.exists("/") else False
                }
            
            return (
                json.dumps(root_paths),
                json.dumps(root_info, indent=2),
                json.dumps(system_info)
            )
            
        except Exception as e:
            print(f"获取系统根目录失败: {e}")
            return ("[]", "{}", "{}")
    
    def _get_drive_type_name(self, drive_type):
        """获取Windows驱动器类型名称"""
        drive_types = {
            0: "DRIVE_UNKNOWN",
            1: "DRIVE_NO_ROOT_DIR",
            2: "DRIVE_REMOVABLE",
            3: "DRIVE_FIXED",
            4: "DRIVE_REMOTE",
            5: "DRIVE_CDROM",
            6: "DRIVE_RAMDISK"
        }
        return drive_types.get(drive_type, f"UNKNOWN_{drive_type}")

class DriveInfoNode:
    """驱动器信息节点 - 获取指定驱动器的详细信息"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "drive_path": ("STRING", {"default": "", "multiline": False}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("drive_info", "free_space", "total_space")
    FUNCTION = "get_drive_info"
    CATEGORY = "AI训练/path_manager"

    def get_drive_info(self, drive_path):
        """获取驱动器详细信息"""
        try:
            if not drive_path or not os.path.exists(drive_path):
                return ("{}", "0", "0")
            
            # 获取磁盘使用情况
            try:
                statvfs = os.statvfs(drive_path)
                block_size = statvfs.f_frsize
                total_blocks = statvfs.f_blocks
                free_blocks = statvfs.f_bavail
                
                total_space = total_blocks * block_size
                free_space = free_blocks * block_size
                
                total_gb = round(total_space / (1024**3), 2)
                free_gb = round(free_space / (1024**3), 2)
                used_gb = total_gb - free_gb
                usage_percent = round((used_gb / total_gb) * 100, 2) if total_gb > 0 else 0
                
            except Exception as e:
                print(f"获取磁盘空间信息失败: {e}")
                total_gb = free_gb = used_gb = usage_percent = 0
            
            # 构建驱动器信息
            drive_info = {
                "path": drive_path,
                "exists": True,
                "accessible": os.access(drive_path, os.R_OK),
                "writable": os.access(drive_path, os.W_OK),
                "is_file": os.path.isfile(drive_path),
                "is_dir": os.path.isdir(drive_path),
                "is_absolute": os.path.isabs(drive_path),
                "space_info": {
                    "total_gb": total_gb,
                    "free_gb": free_gb,
                    "used_gb": used_gb,
                    "usage_percent": usage_percent
                }
            }
            
            # 在Windows上获取额外信息
            if platform.system().lower() == "windows":
                try:
                    from ctypes import windll
                    drive_type = windll.kernel32.GetDriveTypeW(drive_path)
                    drive_info["drive_type"] = self._get_drive_type_name(drive_type)
                    drive_info["drive_type_code"] = drive_type
                except Exception as e:
                    print(f"获取Windows驱动器类型失败: {e}")
            
            return (
                json.dumps(drive_info, indent=2),
                str(free_gb),
                str(total_gb)
            )
            
        except Exception as e:
            print(f"获取驱动器信息失败: {e}")
            return ("{}", "0", "0")
    
    def _get_drive_type_name(self, drive_type):
        """获取Windows驱动器类型名称"""
        drive_types = {
            0: "DRIVE_UNKNOWN",
            1: "DRIVE_NO_ROOT_DIR",
            2: "DRIVE_REMOVABLE",
            3: "DRIVE_FIXED",
            4: "DRIVE_REMOTE",
            5: "DRIVE_CDROM",
            6: "DRIVE_RAMDISK"
        }
        return drive_types.get(drive_type, f"UNKNOWN_{drive_type}")

class PathNavigatorNode:
    """路径导航器节点 - 从根目录开始导航"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "base_path": ("STRING", {"default": "", "multiline": False}),
                "max_depth": ("INT", {"default": 2, "min": 1, "max": 5}),
                "show_hidden": ("BOOLEAN", {"default": False}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("navigation_tree", "current_path", "breadcrumb")
    FUNCTION = "navigate_path"
    CATEGORY = "AI训练/path_manager"

    def navigate_path(self, base_path, max_depth, show_hidden):
        """导航路径并返回树形结构"""
        try:
            if not base_path:
                # 如果没有指定路径，使用当前工作目录
                base_path = os.getcwd()
            
            if not os.path.exists(base_path):
                return ("{}", base_path, "路径不存在")
            
            # 构建导航树
            navigation_tree = self._build_navigation_tree(base_path, max_depth, show_hidden, 0)
            
            # 构建面包屑导航
            breadcrumb = self._build_breadcrumb(base_path)
            
            return (
                json.dumps(navigation_tree, indent=2),
                base_path,
                " > ".join(breadcrumb)
            )
            
        except Exception as e:
            print(f"路径导航失败: {e}")
            return ("{}", base_path, f"错误: {str(e)}")
    
    def _build_navigation_tree(self, path, max_depth, show_hidden, current_depth):
        """构建导航树"""
        if current_depth > max_depth:
            return {}
        
        tree = {}
        
        try:
            for item in os.listdir(path):
                # 跳过隐藏文件（如果未启用显示隐藏文件）
                if not show_hidden and item.startswith('.'):
                    continue
                
                item_path = os.path.join(path, item)
                
                if os.path.isdir(item_path):
                    # 递归构建子目录
                    subtree = self._build_navigation_tree(item_path, max_depth, show_hidden, current_depth + 1)
                    tree[item] = {
                        "type": "folder",
                        "path": item_path,
                        "children": subtree
                    }
                else:
                    # 文件信息
                    try:
                        size = os.path.getsize(item_path)
                        tree[item] = {
                            "type": "file",
                            "path": item_path,
                            "size": size,
                            "size_mb": round(size / (1024 * 1024), 2)
                        }
                    except Exception as e:
                        tree[item] = {
                            "type": "file",
                            "path": item_path,
                            "error": str(e)
                        }
        except PermissionError:
            tree["[权限不足]"] = {"type": "error", "message": "无法访问此目录"}
        except Exception as e:
            tree[f"[错误]"] = {"type": "error", "message": str(e)}
        
        return tree
    
    def _build_breadcrumb(self, path):
        """构建面包屑导航"""
        try:
            if platform.system().lower() == "windows":
                # Windows路径处理
                parts = path.split(os.sep)
                if len(parts) > 1 and parts[1] == "":  # 绝对路径
                    breadcrumb = [parts[0] + "\\"]  # 盘符
                    breadcrumb.extend(parts[2:])
                else:
                    breadcrumb = parts
            else:
                # Unix/Linux/macOS路径处理
                breadcrumb = path.split(os.sep)
                if breadcrumb[0] == "":  # 绝对路径
                    breadcrumb[0] = "/"
            
            # 过滤空字符串
            breadcrumb = [part for part in breadcrumb if part]
            return breadcrumb
            
        except Exception as e:
            return [f"错误: {str(e)}"]

# 节点映射
NODE_CLASS_MAPPINGS = {
    "SystemRootsNode": SystemRootsNode,
    "DriveInfoNode": DriveInfoNode,
    "PathNavigatorNode": PathNavigatorNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SystemRootsNode": "系统根目录",
    "DriveInfoNode": "驱动器信息",
    "PathNavigatorNode": "路径导航器",
} 